import math
from functools import partial

import numpy as np
import tensorflow as tf
from tensorflow.keras import backend as K

# 定义带有类别权重、Beta系数和平滑项的Dice损失与交叉熵损失结合的函数
def dice_loss_with_CE(cls_weights, beta=1, smooth=1e-5):
    # 将类别权重重塑为四维张量，以便在损失计算中使用
    cls_weights = np.reshape(cls_weights, [1, 1, 1, -1])
    # 定义内部函数，用于计算损失
    def _dice_loss_with_CE(y_true, y_pred):
        # 将预测概率限制在一个小范围内，以避免数值问题
        y_pred = K.clip(y_pred, K.epsilon(), 1.0 - K.epsilon())
        # 计算交叉熵损失，只考虑除了背景类别之外的所有类别
        CE_loss = - y_true[...,:-1] * K.log(y_pred) * cls_weights
        # 计算交叉熵损失的平均值
        CE_loss = K.mean(K.sum(CE_loss, axis=-1))
        # 计算真正例（TP）、假正例（FP）和假负例（FN）
        tp = K.sum(y_true[...,:-1] * y_pred, axis=[0,1,2])
        fp = K.sum(y_pred, axis=[0,1,2]) - tp
        fn = K.sum(y_true[...,:-1], axis=[0,1,2]) - tp
        # 计算Dice系数
        score = ((1 + beta ** 2) * tp + smooth) / ((1 + beta ** 2) * tp + beta ** 2 * fn + fp + smooth)
        # 计算Dice损失
        dice_loss = 1 - score
        # 返回总损失，即交叉熵损失与Dice损失的和
        return CE_loss + dice_loss
    # 返回定义好的损失函数
    return _dice_loss_with_CE

# 定义纯交叉熵损失函数
def CE(cls_weights):
    # 类别权重的处理与上面相同
    cls_weights = np.reshape(cls_weights, [1, 1, 1, -1])
    # 定义内部函数，用于计算交叉熵损失
    def _CE(y_true, y_pred):
        # 预测概率的限制与上面相同
        y_pred = K.clip(y_pred, K.epsilon(), 1.0 - K.epsilon())
        # 计算交叉熵损失，只考虑除了背景类别之外的所有类别
        CE_loss = - y_true[...,:-1] * K.log(y_pred) * cls_weights
        # 计算交叉熵损失的平均值
        CE_loss = K.mean(K.sum(CE_loss, axis=-1))
        # 返回交叉熵损失
        return CE_loss
    # 返回定义好的损失函数
    return _CE

# 定义带有类别权重、Beta系数、平滑项、Alpha和Gamma参数的Dice损失与Focal损失结合的函数
def dice_loss_with_Focal_Loss(cls_weights, beta=1, smooth=1e-5, alpha=0.5, gamma=2):
    # 类别权重的处理与上面相同
    cls_weights = np.reshape(cls_weights, [1, 1, 1, -1])
    # 定义内部函数，用于计算损失
    def _dice_loss_with_Focal_Loss(y_true, y_pred):
        # 预测概率的限制与上面相同
        y_pred = K.clip(y_pred, K.epsilon(), 1.0 - K.epsilon())
        # 计算未归一化的对数损失
        logpt = - y_true[...,:-1] * K.log(y_pred) * cls_weights
        # 对所有类别的对数损失求和
        logpt = - K.sum(logpt, axis=-1)
        # 计算预测概率
        pt = tf.exp(logpt)
        # 如果提供了Alpha参数，按类别权重调整对数损失
        if alpha is not None:
            logpt *= alpha
        # 计算Focal损失
        CE_loss = -((1 - pt) ** gamma) * logpt
        # 计算Focal损失的平均值
        CE_loss = K.mean(CE_loss)
        # Dice损失的计算与上面相同
        # tp, fp, fn = K.sum(y_true[...,:-1] * y_pred, axis=[0,1,2]), K.sum(y_pred, axis=[0,1,2]) - tp, K.sum(y_true[...,:-1], axis=[0,1,2]) - tp
        tp = K.sum(y_true[..., :-1] * y_pred, axis=[0, 1, 2])
        fp = K.sum(y_pred, axis=[0, 1, 2]) - tp
        fn = K.sum(y_true[..., :-1], axis=[0, 1, 2]) - tp

        score = ((1 + beta ** 2) * tp + smooth) / ((1 + beta ** 2) * tp + beta ** 2 * fn + fp + smooth)
        dice_loss = 1 - score
        # 返回总损失，即Focal损失与Dice损失的和
        return CE_loss + dice_loss
    # 返回定义好的损失函数
    return _dice_loss_with_Focal_Loss

# 定义Focal损失函数
def Focal_Loss(cls_weights, alpha=0.5, gamma=2):
    # 类别权重的处理与上面相同
    cls_weights = np.reshape(cls_weights, [1, 1, 1, -1])
    # 定义内部函数，用于计算Focal损失
    def _Focal_Loss(y_true, y_pred):
        # 预测概率的限制与上面相同
        y_pred = K.clip(y_pred, K.epsilon(), 1.0 - K.epsilon())
        # 计算未归一化的对数损失
        logpt = - y_true[...,:-1] * K.log(y_pred) * cls_weights
        # 对所有类别的对数损失求和
        logpt = - K.sum(logpt, axis=-1)
        # 计算预测概率
        pt = tf.exp(logpt)
        # 如果提供了Alpha参数，按类别权重调整对数损失
        if alpha is not None:
            logpt *= alpha
        # 计算Focal损失
        CE_loss = -((1 - pt) ** gamma) * logpt
        # 计算Focal损失的平均值
        CE_loss = K.mean(CE_loss)
        # 返回Focal损失
        return CE_loss
    # 返回定义好的损失函数
    return _Focal_Loss

# 定义一个学习率调度器生成函数
# 这个get_lr_scheduler函数的目的是根据不同的学习率衰减类型（余弦衰减或步长衰减）来创建一个学习率调度器。
# 这个调度器可以在训练过程中动态调整学习率。函数内部定义了两个学习率调度器：yolox_warm_cos_lr用于余弦衰减学习率，
# step_lr用于步长衰减学习率。然后根据传入的lr_decay_type参数选择使用哪一个调度器。

import math
def get_lr_scheduler(lr_decay_type, lr, min_lr, total_epochs):

    def lr(epoch,lr):
        epoch=epoch +1
        current_lr = lr
        warmup_total_epochs = int(total_epochs * 0.1)  # 假设前10%的epochs为warmup
        if warmup_total_epochs == 0:
            warmup_total_epochs = 1  # 避免除以零，至少设置为1
        print(f"Epoch: {epoch}, Warmup Total Epochs: {warmup_total_epochs},total_epochs:{total_epochs}")
        if lr_decay_type == 'cos':
            warmup_total_epochs = int(total_epochs * 0.1)
            if epoch <= warmup_total_epochs:
                current_lr = (lr - 0) * (epoch / float(warmup_total_epochs))
            else:
                decay = 0.5 * (1 + math.cos(math.pi * (epoch - warmup_total_epochs) / (total_epochs - warmup_total_epochs)))
                current_lr = lr * max(decay, 0.01)

        elif lr_decay_type == 'step':
            step_size = int(total_epochs / 10)
            if epoch % step_size == 0:
                current_lr *= 0.1

        current_lr = max(current_lr, min_lr)
        return current_lr

    return lr




#
#
# def get_lr_scheduler(lr_decay_type, lr, min_lr, total_iters, warmup_iters_ratio=0.05, warmup_lr_ratio=0.1, no_aug_iter_ratio=0.05, step_num=10):
#     # 定义余弦衰减学习率调度器
#
#
#     def yolox_warm_cos_lr(lr, min_lr, total_iters, warmup_total_iters, warmup_lr_start, no_aug_iter, iters):
#         if iters <= warmup_total_iters:
#             # 在预热阶段，学习率从warmup_lr_start线性增加到lr
#             lr = (lr - warmup_lr_start) * pow(iters / float(warmup_total_iters), 2) + warmup_lr_start
#         elif iters >= total_iters - no_aug_iter:
#             # 在训练的最后no_aug_iter迭代中，学习率保持为min_lr
#             lr = min_lr
#         else:
#             # 在训练的其余阶段，学习率按照余弦函数衰减
#             lr = min_lr + 0.5 * (lr - min_lr) * (1.0 + math.cos(math.pi * (iters - warmup_total_iters) / (total_iters - warmup_total_iters - no_aug_iter)))
#         return lr
#
#     # 定义步长衰减学习率调度器
#     def step_lr(lr, decay_rate, step_size, iters):
#         # 计算迭代次数对应的衰减次数
#         n = iters // step_size
#         # 计算当前迭代的学习率
#         out_lr = lr * decay_rate ** n
#         return out_lr
#
#     # 根据lr_decay_type选择学习率调度器
#     if lr_decay_type == "cos":
#         # 预热迭代次数和学习率的计算
#         warmup_total_iters = min(max(warmup_iters_ratio * total_iters, 1), 3)
#         warmup_lr_start = max(warmup_lr_ratio * lr, 1e-6)
#         no_aug_iter = min(max(no_aug_iter_ratio * total_iters, 1), 15)
#         # 使用partial函数创建余弦衰减学习率调度器
#         func = partial(yolox_warm_cos_lr, lr, min_lr, total_iters, warmup_total_iters, warmup_lr_start, no_aug_iter)
#
#     else:
#         # 步长衰减学习率的计算
#         decay_rate = (min_lr / lr) ** (1 / (step_num - 1))
#         step_size = total_iters / step_num
#         # 使用partial函数创建步长衰减学习